def load_dataset():
    """创建示例数据集"""
    return [
        ['牛奶', '面包', '尿布'],
        ['可乐', '面包', '尿布', '啤酒'],
        ['牛奶', '尿布', '啤酒', '鸡蛋'],
        ['面包', '牛奶', '尿布', '啤酒'],
        ['面包', '牛奶', '尿布', '可乐']
    ]


def create_c1(dataset):
    """创建初始候选项集"""
    c1 = []
    for transaction in dataset:
        for item in transaction:
            if [item] not in c1:
                c1.append([item])
    c1.sort()
    return list(map(frozenset, c1))


def scan_dataset(dataset, candidates, min_support):
    """扫描数据集，计算支持度"""
    item_count = {}
    for transaction in dataset:
        for candidate in candidates:
            if candidate.issubset(transaction):
                item_count[candidate] = item_count.get(candidate, 0) + 1

    num_transactions = float(len(dataset))
    frequent_items = []
    support_data = {}

    for candidate, count in item_count.items():
        support = count / num_transactions
        if support >= min_support:
            frequent_items.append(candidate)
        support_data[candidate] = support

    return frequent_items, support_data


def apriori_gen(frequent_items, k):
    """生成候选项集"""
    candidates = []
    len_frequent = len(frequent_items)

    for i in range(len_frequent):
        for j in range(i + 1, len_frequent):
            l1 = list(frequent_items[i])[:k - 2]
            l2 = list(frequent_items[j])[:k - 2]
            l1.sort()
            l2.sort()

            if l1 == l2:
                candidates.append(frequent_items[i] | frequent_items[j])

    return candidates


def apriori(dataset, min_support=0.5):
    """Apriori算法主函数"""
    c1 = create_c1(dataset)
    dataset = list(map(set, dataset))

    frequent_1, support_data = scan_dataset(dataset, c1, min_support)
    frequent_items = [frequent_1]
    k = 2

    while len(frequent_items[k - 2]) > 0:
        candidates = apriori_gen(frequent_items[k - 2], k)
        frequent_k, support_k = scan_dataset(dataset, candidates, min_support)

        support_data.update(support_k)
        frequent_items.append(frequent_k)
        k += 1

    return frequent_items, support_data


# 测试代码
if __name__ == "__main__":
    dataset = load_dataset()
    min_support = 0.4

    frequent_items, support_data = apriori(dataset, min_support)

    print("频繁项集:")
    for i, items in enumerate(frequent_items):
        if items:
            print(f"长度 {i + 1}: {[list(item) for item in items]}")

    print("\n支持度:")
    for item, support in support_data.items():
        if support >= min_support:
            print(f"{list(item)}: {support:.2f}")